#!/bin/bash

# Copyright 2017 Xingyu Na
# Apache 2.0

. ./path.sh || exit 1;
#需要修改path
if [ $# != 2 ]; then
  echo "Usage: $0 <audio-path> <text-path>"
  echo " $0 ~/autodl-tmp/wenet/data/sentence_data/wav ~/autodl-tmp/wenet/data/sentence_data/transcript"
  exit 1;
fi

sentence_audio_dir=$1
transcript_dir=$2

train_dir=data/local/train
dev_dir=data/local/dev
test_dir=data/local/test
tmp_dir=data/local/tmp

mkdir -p $train_dir
mkdir -p $dev_dir
mkdir -p $test_dir
mkdir -p $tmp_dir

# 检查目录是否存在
if [ ! -d $sentence_audio_dir ] || [ ! -d $transcript_dir ]; then
  echo "Error: $0 requires two directory arguments"
  exit 1;
fi

# 1. 处理音频文件 - 按数据集划分
for dataset in train dev test; do
  # 创建每个数据集的wav.flist
  find $sentence_audio_dir/$dataset -iname "*.wav" > $tmp_dir/${dataset}_wav.flist
  sed "s|$sentence_audio_dir/$dataset/||" $tmp_dir/${dataset}_wav.flist > $tmp_dir/${dataset}_wav.flist.tmp
  mv $tmp_dir/${dataset}_wav.flist.tmp $tmp_dir/${dataset}_wav.flist
  
  # 创建utt.list (使用文件名作为utterance ID)
  awk -F/ '{print $NF}' $tmp_dir/${dataset}_wav.flist | sed 's/\.wav//' > $tmp_dir/${dataset}_utt.list
done

# 2. 处理转录文件 - 合并所有转录文件
# 先清洗转录文件中的括号内容
python ~/autodl-tmp/wenet/wenet/examples/senior_talk/s0/local/clean_data.py $transcript_dir

# 合并所有转录文件
cat $transcript_dir/*.txt > $tmp_dir/all_transcripts.txt

# 3. 为每个数据集创建wav.scp和text文件
for dataset in train dev test; do
  # 创建wav.scp
  #paste -d' ' $tmp_dir/${dataset}_utt.list $tmp_dir/${dataset}_wav.flist > $tmp_dir/${dataset}_wav.scp_all
  paste -d' ' $tmp_dir/${dataset}_utt.list <(sed "s|^|$sentence_audio_dir/$dataset/|" $tmp_dir/${dataset}_wav.flist) > $tmp_dir/${dataset}_wav.scp_all
  # 创建text文件
  while read utt_id; do
    # 查找对应的转录文本
    transcript_line=$(grep -m1 "$utt_id" $tmp_dir/all_transcripts.txt || echo "")
    if [ -n "$transcript_line" ]; then
      # 提取文本部分 (假设格式: utterance_id 文本内容)
      transcript_text=$(echo "$transcript_line" | awk '{$1=""; print $0}' | sed 's/^ //')
      echo "$utt_id $transcript_text" >> $tmp_dir/${dataset}_text.tmp
    else
      echo "Warning: No transcript found for $utt_id"
    fi
  done < $tmp_dir/${dataset}_utt.list
  
  # 排序并去重
  sort -u $tmp_dir/${dataset}_text.tmp > $tmp_dir/${dataset}_text
  rm $tmp_dir/${dataset}_text.tmp
  
  # 创建Kaldi数据目录
  mkdir -p data/$dataset
  
  # 复制最终文件
  cp $tmp_dir/${dataset}_wav.scp_all data/$dataset/wav.scp
  cp $tmp_dir/${dataset}_text data/$dataset/text
  
  # 创建utt2spk (使用文件名作为speaker ID)
  awk '{print $1, $1}' data/$dataset/wav.scp > data/$dataset/utt2spk
  
  # 创建spk2utt
  ~/autodl-tmp/wenet/wenet/tools/utt2spk_to_spk2utt.pl data/$dataset/utt2spk > data/$dataset/spk2utt
done

# 4. 创建全局CMVN
# train_set=train
# tools/compute_cmvn_stats.py --num_workers 8 \
#   --in_scp data/${train_set}/wav.scp \
#   --out_cmvn data/$train_set/global_cmvn

echo "$0: SeniorTalk data preparation succeeded"
exit 0;